from Settings import OutBaseDir

from insightface.app import FaceAnalysis

import cv2
import numpy as np

from typing import List
from insightface.app.common import Face

class PointFace:
    def __init__(self):
        self.__pointFaceEngine = FaceAnalysis(allowed_modules=['detection', 'landmark_2d_106'])
        self.__pointFaceEngine.prepare(ctx_id=0, det_size=(640, 640))

    def loadFaceImage(self, image_path):
        img = cv2.imread(image_path)
        faces_data: List[Face] = self.__pointFaceEngine.get(img)

        return img, faces_data

    def drawPointFaceImage(self, img, tmp_faces_data: List[Face]):
        tim = img.copy()
        # 设置轮廓点的颜色
        color = (0, 0, 255)
        for face in tmp_faces_data:
            lmk = face.landmark_2d_106
            lmk = np.round(lmk).astype(np.int32)
            for i in range(lmk.shape[0]):
                p = tuple(lmk[i])

                cv2.circle(tim, p, 1, color, 1, cv2.LINE_AA)
        cv2.imwrite(f'{OutBaseDir}/test_out1.jpg', tim)

if __name__ == '__main__':
    from Settings import BaseTestDataCubePath

    pointFace = PointFace()
    imagePath = f'{BaseTestDataCubePath}/1.jpg'
    print(imagePath)
    original_img, faces_data = pointFace.loadFaceImage(imagePath)
    pointFace.drawPointFaceImage(original_img, faces_data)


